"""
修复版三维建筑疏散仿真代码
包含：数值稳定性修复、楼梯容量同步控制、异常处理机制
"""

import numpy as np
import heapq
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation

# 论文参数常量
STAIR_CAPACITY = 15        # 楼梯最大容量
BASE_SPEED_MALE = 1.34     # 男性基础速度系数
BASE_SPEED_FEMALE = 1.25   # 女性基础速度系数
STAIR_UP_COST = 3.2        # 上楼梯体力消耗
STAIR_DOWN_COST = 1.8      # 下楼梯体力消耗
FLOOR_HEIGHT = 3.0         # 层高（米）
MAX_FLOORS = 20            # 建筑总层数

class Building3D:
    def __init__(self, floors=MAX_FLOORS, size=50):
        self.floors = floors
        self.size = size
        self.stair_locations = {}
        
        # 每5层设置双向楼梯
        for f in range(0, floors, 5):
            mid = size // 2
            connections = []
            if f-5 >= 0: connections.append(f-5)
            if f+5 < floors: connections.append(f+5)
            
            self.stair_locations[f] = {
                'position': (mid, mid),
                'current_people': 0,
                'connections': connections
            }
        
        # 创建楼层网格（0-可通行 1-墙壁 2-楼梯）
        self.grid = np.zeros((floors, size, size))
        for f in self.stair_locations:
            x, y = self.stair_locations[f]['position']
            self.grid[f, x, y] = 2

class Pedestrian:
    def __init__(self, floor, position, age=35, gender='male'):
        self.floor = floor
        self.position = np.array(position, dtype=np.float32)
        self.age = np.clip(age, 20, 60)
        self.gender = gender
        self.max_stamina = self._calc_max_stamina()
        self.stamina = self.max_stamina
        self.speed = self._calc_speed()
        self.in_stair = False
    
    def _calc_max_stamina(self):
        base = 3000 if self.gender == 'male' else 2500
        return base * (0.85 ** ((self.age - 25) / 5))
    
    def _calc_speed(self):
        if self.gender == 'male':
            base = BASE_SPEED_MALE - 0.0065*self.age
        else:
            base = BASE_SPEED_FEMALE - 0.0072*self.age
        stamina_factor = 0.4 + 0.6*(self.stamina/self.max_stamina)**0.8
        return base * stamina_factor
    
    def update_stamina(self, movement_type):
        cost = {'walk':0.5, 'stair_up':3.2, 'stair_down':1.8, 'wait':0.1}[movement_type]
        self.stamina = max(0, self.stamina - cost)
        self.speed = self._calc_speed()

class EvacuationSimulator:
    def __init__(self, building, people):
        self.building = building
        self.people = people
        self.time_step = 0
        self.dynamic_field = self._generate_dynamic_field()
    
    def _generate_dynamic_field(self):
        field = np.full_like(self.building.grid, np.inf, dtype=np.float32)
        exit_floor = 0
        exit_pos = self.building.stair_locations[exit_floor]['position']
        heap = [(0, exit_floor, exit_pos[0], exit_pos[1])]
        field[exit_floor, exit_pos[0], exit_pos[1]] = 0
        
        directions = [(-1,-1), (-1,0), (-1,1), (0,-1), (0,1), (1,-1), (1,0), (1,1)]
        
        while heap:
            current_cost, floor, x, y = heapq.heappop(heap)
            
            # 处理楼梯连接
            if self.building.grid[floor, x, y] == 2:
                for cf in self.building.stair_locations[floor]['connections']:
                    if cf >= self.building.floors: continue
                    stair_pos = self.building.stair_locations[cf]['position']
                    new_cost = current_cost + 5  # 楼梯额外成本
                    if new_cost < field[cf, stair_pos[0], stair_pos[1]]:
                        field[cf, stair_pos[0], stair_pos[1]] = new_cost
                        heapq.heappush(heap, (new_cost, cf, stair_pos[0], stair_pos[1]))
            
            # 处理平面移动
            for dx, dy in directions:
                nx, ny = x+dx, y+dy
                if 0<=nx<self.building.size and 0<=ny<self.building.size:
                    if self.building.grid[floor, nx, ny] != 1:
                        move_cost = np.sqrt(dx**2 + dy**2)
                        if (current_cost + move_cost) < field[floor, nx, ny]:
                            field[floor, nx, ny] = current_cost + move_cost
                            heapq.heappush(heap, (field[floor, nx, ny], floor, nx, ny))
        
        # 处理未连接区域
        field[np.isinf(field)] = 1000  # 设置极大值避免计算错误
        return field
    
    def update(self):
        for person in self.people:
            if person.floor == 0 and np.linalg.norm(person.position - self.building.stair_locations[0]['position']) < 1:
                continue
            
            current_floor = int(person.floor)
            current_x, current_y = person.position.astype(int)
            possible_moves = []
            
            # 平面移动
            for dx in [-1,0,1]:
                for dy in [-1,0,1]:
                    if dx==0 and dy==0: continue
                    nx, ny = current_x+dx, current_y+dy
                    if 0<=nx<self.building.size and 0<=ny<self.building.size:
                        if self.building.grid[current_floor, nx, ny] != 1:
                            cost = self.dynamic_field[current_floor, nx, ny]
                            possible_moves.append( (current_floor, nx, ny, cost, 'walk') )
            
            # 楼梯移动（严格容量检查）
            if self.building.grid[current_floor, current_x, current_y] == 2:
                stair_info = self.building.stair_locations[current_floor]
                # 上楼梯
                if any([cf for cf in stair_info['connections'] if cf > current_floor]):
                    target_floor = min([cf for cf in stair_info['connections'] if cf > current_floor], default=None)
                    if target_floor is not None:
                        target_stair = self.building.stair_locations[target_floor]
                        if (stair_info['current_people'] < STAIR_CAPACITY and 
                            target_stair['current_people'] < STAIR_CAPACITY):
                            target_pos = target_stair['position']
                            cost = self.dynamic_field[target_floor, target_pos[0], target_pos[1]] + 5
                            possible_moves.append( (target_floor, target_pos[0], target_pos[1], cost, 'stair_up') )
                # 下楼梯
                if any([cf for cf in stair_info['connections'] if cf < current_floor]):
                    target_floor = max([cf for cf in stair_info['connections'] if cf < current_floor], default=None)
                    if target_floor is not None:
                        target_stair = self.building.stair_locations[target_floor]
                        if (stair_info['current_people'] < STAIR_CAPACITY and 
                            target_stair['current_people'] < STAIR_CAPACITY):
                            target_pos = target_stair['position']
                            cost = self.dynamic_field[target_floor, target_pos[0], target_pos[1]] + 3
                            possible_moves.append( (target_floor, target_pos[0], target_pos[1], cost, 'stair_down') )
            
            if not possible_moves:
                person.update_stamina('wait')
                continue
            
            # 概率计算（增强稳定性）
            try:
                possible_moves.sort(key=lambda x: x[3])
                best_moves = possible_moves[:3]
                
                costs = np.array([m[3] for m in best_moves])
                min_cost = np.min(costs)
                adjusted_costs = costs - min_cost
                probabilities = np.exp(-0.5 * adjusted_costs)
                sum_prob = probabilities.sum()
                
                if sum_prob < 1e-10 or np.any(np.isnan(probabilities)):
                    probabilities = np.ones(len(best_moves)) / len(best_moves)
                else:
                    probabilities /= sum_prob
                
                choice = np.random.choice(len(best_moves), p=probabilities)
                target_floor, tx, ty, _, move_type = best_moves[choice]
                
                # 楼梯状态更新
                if move_type.startswith('stair'):
                    prev_stair = self.building.stair_locations[current_floor]
                    next_stair = self.building.stair_locations[target_floor]
                    
                    if person.in_stair:
                        prev_stair['current_people'] = max(0, prev_stair['current_people']-1)
                    
                    if next_stair['current_people'] < STAIR_CAPACITY:
                        next_stair['current_people'] += 1
                        person.in_stair = True
                    else:
                        continue  # 楼梯已满，取消移动
                
                # 位置更新
                person.floor = target_floor
                person.position = np.array([tx, ty])
                person.update_stamina(move_type)
                
            except Exception as e:
                person.update_stamina('wait')

class Visualization3D:
    def __init__(self, simulator):
        self.sim = simulator
        self.fig = plt.figure(figsize=(12, 9))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.setup_axes()
        
    def setup_axes(self):
        self.ax.set_xlim(0, self.sim.building.size)
        self.ax.set_ylim(0, self.sim.building.size)
        self.ax.set_zlim(0, self.sim.building.floors*FLOOR_HEIGHT)
        self.ax.set_xlabel('X Position')
        self.ax.set_ylabel('Y Position')
        self.ax.set_zlabel('Floor Height')
        self.ax.view_init(elev=30, azim=45)
        
    def update_plot(self, frame):
        self.sim.update()
        self.ax.clear()
        self.setup_axes()
        
        # 绘制楼梯
        for f in self.sim.building.stair_locations:
            x, y = self.sim.building.stair_locations[f]['position']
            self.ax.scatter(x, y, f*FLOOR_HEIGHT+0.5, c='red', marker='s', s=50, alpha=0.5)
        
        # 绘制人群
        colors, positions = [], []
        for p in self.sim.people:
            x, y = p.position
            z = p.floor * FLOOR_HEIGHT + 0.5
            positions.append([x, y, z])
            colors.append(p.stamina / p.max_stamina)
        
        self.ax.scatter(
            xs=[p[0] for p in positions],
            ys=[p[1] for p in positions],
            zs=[p[2] for p in positions],
            c=colors, cmap='RdYlGn', alpha=0.7, 
            vmin=0, vmax=1, s=20
        )
        return self.ax,
    
    def animate(self):
        ani = FuncAnimation(self.fig, self.update_plot, frames=200, interval=50, blit=False)
        plt.show()

if __name__ == "__main__":
    # 初始化建筑和人群
    building = Building3D(floors=MAX_FLOORS)
    people = []
    
    for floor in range(MAX_FLOORS):
        n_people = np.random.randint(50, 100)
        for _ in range(n_people):
            x = np.random.uniform(0, building.size)
            y = np.random.uniform(0, building.size)
            age = np.random.normal(35, 10)
            gender = 'male' if np.random.rand() > 0.5 else 'female'
            people.append(Pedestrian(floor, (x, y), age, gender))
    
    # 启动仿真
    simulator = EvacuationSimulator(building, people)
    vis = Visualization3D(simulator)
    vis.animate()
